package com.inspur.cmp.jov.keycloak.tools;

import cn.hutool.core.codec.Base64Decoder;
import cn.hutool.http.HttpUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.fasterxml.jackson.databind.type.CollectionType;
import com.inspur.cmp.jov.keycloak.tools.model.JWS;

import javax.annotation.Nonnull;
import java.math.BigInteger;
import java.security.KeyFactory;
import java.security.NoSuchAlgorithmException;
import java.security.PublicKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.RSAPublicKeySpec;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

public class CertManager {

    public static final String JWKS_ENDPOINT = "/auth/realms/%s/protocol/openid-connect/certs";

    /**
     * 缓存时长(millis)
     */
    public static final long CERT_CACHE_TIME = 10 * 60 * 1000;

    private Map<String, RSAPublicKey> publicKeyMap = new HashMap<>();
    private long timestamp;

    private ObjectMapper objectMapper;

    private String authServer;

    /**
     * 证书完整地址
     */
    private String jwksUri;

    public CertManager(@Nonnull String authServer, @Nonnull String realm) {
        Objects.requireNonNull(authServer);
        Objects.requireNonNull(realm);

        this.authServer = authServer;
        this.jwksUri = String.format(authServer + JWKS_ENDPOINT, realm);

        this.objectMapper = new ObjectMapper();
        objectMapper.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
    }



    public RSAPublicKey getPublicKey(@Nonnull String kid){
        RSAPublicKey result;
        if (timestamp + CERT_CACHE_TIME > System.currentTimeMillis()) {
            result =  publicKeyMap.get(kid);
        }else {
            getCertsFromAuthServer();
            result =  publicKeyMap.get(kid);
        }
        if (result == null) {
            throw new RuntimeException("无法识别的kid");
        }
        return result;
    }

    /**
     * 目前IAM默认使用RSA+sha256算法进行token签名，暂时忽略其他情况
     */
    private synchronized void getCertsFromAuthServer() {

        if (timestamp + CERT_CACHE_TIME > System.currentTimeMillis()) {
            return;
        }
        //all certs
        String json = HttpUtil.get(jwksUri);
        try {
            ObjectNode root = (ObjectNode) objectMapper.readTree(json);
            JsonNode keys =  root.path("keys");
            CollectionType type = objectMapper.getTypeFactory().constructCollectionType(List.class, JWS.class);
            List<JWS> list = objectMapper.treeToValue(keys, type);
            if (list != null && !list.isEmpty()) {
                KeyFactory keyFactory = KeyFactory.getInstance("RSA");
                for (JWS jws : list) {
                    RSAPublicKeySpec rsaPublicKeySpec = new RSAPublicKeySpec(
                            new BigInteger(1, Base64Decoder.decode(jws.getN())),
                            new BigInteger(1, Base64Decoder.decode(jws.getE()))
                    );
                    if (jws.getUse().equals("sig")) {

                        PublicKey publicKey = keyFactory.generatePublic(rsaPublicKeySpec);
                        publicKeyMap.put(jws.getKid(), (RSAPublicKey) publicKey);
                        timestamp = System.currentTimeMillis();
                    }
                }
            }
        } catch (JsonProcessingException e) {
            throw new RuntimeException(e);
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        } catch (InvalidKeySpecException e) {
            throw new RuntimeException(e);
        }
    }
}
